import numpy as np
from catboost import CatBoostClassifier
import pandas as pd
from collections import defaultdict
import os
from utils.parser import parse_args

if __name__ == '__main__':
    args = parse_args()

    id = str(args.run_id)
    file_id = '2'

    data = pd.read_table('../../data/three_week_data.txt', sep='\t')
    test = pd.read_table('../../data/orders_test_spu.txt', sep='\t')
    ori_train = data

    # 日期转换
    dt_list = list(range(20210607, 20210631)) + [20210701, 20210702]
    dt_trans_dict = {}
    num = 0
    for dt in dt_list:
        dt_trans_dict[dt] = num
        num += 1

    if os.path.exists('data/submit-train_'+file_id+'.csv'): # 如果检测到没有训练数据，则生成
        pass
    else:
        train = pd.DataFrame(ori_train, columns=['user_id', 'aor_id', 'ord_period_name', 'dt', 'wm_food_spu_id'])
        train['dt'] = train['dt'].apply(lambda x : dt_trans_dict[x])
        recall_items = list(pd.read_table('../solution_1/data/submit_KGIN_data/recall_items_2w.txt', header=None)[0])
        user_spu = defaultdict(list)
        for line in open('../solution_1/data/submit_KGIN_data/train.txt'):
            line = line.strip()
            lines = line.split(' ')
            lines = [int(a) for a in lines]
            user_spu[lines[0]].extend(list(set(lines[1:])))
        neg_train = train
        neg_list = []
        for user_id in neg_train['user_id']:
            while True:
                neg_item = np.random.randint(low=0, high=len(recall_items), size=1)[0]
                neg_item = recall_items[neg_item]
                if neg_item not in user_spu[user_id]:
                    break
            neg_list.append(neg_item)
        neg_spu = pd.DataFrame(neg_list, columns=['wm_food_spu_id'])
        neg_train = neg_train.drop('wm_food_spu_id', axis=1)
        neg_train = pd.concat([neg_train, neg_spu], axis=1)
        new_train = pd.concat([train, neg_train])
        new_train.to_csv('data/submit-train_'+file_id+'.csv', index=False)

    user_bprmf_recall_spu = defaultdict(list)
    for line in open('../solution_1/result/3_top_50_submit_KGIN_data_KGIN_RI_test.txt'): # two_top_50_KGIN_RI_5_50.txt  two_top_50_KGIN_3_50_recall_2w.txt
        lines = line.strip().split('\t')
        user_bprmf_recall_spu[int(lines[0])].extend([int(a) for a in lines[1:]])

    unique_test_ori = test[-test.duplicated(['wm_order_id'])]
    if os.path.exists('data/submit-rank_'+file_id+'.csv'): # 如果检测到没有测试数据，则生成
        pass
    else:
        unique_test = pd.DataFrame(unique_test_ori, columns=['user_id', 'aor_id', 'ord_period_name', 'dt', 'wm_food_spu_id'])
        unique_test['dt'] = unique_test['dt'].apply(lambda x : dt_trans_dict[x])
        unique_test_list = list(np.array(unique_test))
        # 这里目前必须这样写，否则可能存在浅拷贝的问题
        rank_list = list()
        for one in unique_test_list:
            for j in user_bprmf_recall_spu[one[0]]:
                rank_list.append(list(one[:-1])+[j])
        rank_spus = pd.DataFrame(rank_list, columns=['user_id', 'aor_id', 'ord_period_name', 'dt', 'wm_food_spu_id'],
                                 dtype=np.int64)
        rank_spus.to_csv('data/submit-rank_'+file_id+'.csv',index=False)

    train = pd.read_csv('data/submit-train_'+file_id+'.csv')
    train_label = pd.DataFrame([1] * (len(train)//2) + [0] * (len(train)//2))

    categorical_features_indices = np.where(train.dtypes != np.float)[0]

    model = CatBoostClassifier(
        custom_metric=['Accuracy'], # Accuracy,AUC
        random_seed=666,
        eval_metric='Accuracy',
        iterations=args.epoch,
        task_type="GPU",
        depth=args.depth,
        loss_function='Logloss',
    )

    model.fit(
        train, train_label,
        cat_features=categorical_features_indices,
        verbose=10,
        early_stopping_rounds=200,
        plot=False
    )

    model.save_model('model/submit-model_'+id+'.dump')

    test_data = pd.read_csv('data/submit-rank_'+file_id+'.csv')

    result = model.predict_proba(test_data)
    print(result[:10])

    with open('result/submit-model_'+id+'.txt','w') as f:
        for r in result[:,1]:
            f.write(str(r))
            f.write('\n')

    catboost_out = pd.read_table('./result/submit-model_' + id + '.txt', header=None)

    index = 0
    recall_num = 50
    # 存放按照订单重排的结果
    order_rank = {}

    # 生成每个订单推荐50个的结果
    with open('./submit/submit-catboost_50_spu_' + id + '.txt', 'w') as f:
        for tup in zip(unique_test_ori['user_id'], unique_test_ori['wm_order_id']):
            spu_score = {}
            for spu in user_bprmf_recall_spu[tup[0]]:
                spu_score[spu] = catboost_out[0][index]
                index += 1
            # 控制召回的数目
            count = 0
            spu_score_c = {}
            for key in spu_score:
                spu_score_c[key] = spu_score[key]
                count += 1
                if count == recall_num:
                    break
            S = sorted(spu_score_c.items(), key=lambda item: item[1], reverse=True)
            #         S = heapq.nlargest(5, spu_score_c, key=spu_score_c.get)

            order_rank[tup[1]] = S
        for tup in test['wm_order_id']:
            f.write(str(tup))

            for s in order_rank[tup]:
                f.write('\t')
                f.write(str(s[0]))
            f.write('\n')

    # 生成提交给网站的结果
    with open('./submit/submit-catboost_5_spu_' + id + '.txt', 'w') as f:
        for tup in test['wm_order_id']:
            f.write(str(tup))
            count = 0
            for s in order_rank[tup]:
                f.write('\t')
                f.write(str(s[0]))
                count += 1
                if count == 5:
                    break
            f.write('\n')
    print('完成训练')
    print(model.get_all_params())